package javaagent.jdbc;

import javassist.*;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;

/**
 * @author ：zzs
 * @version : 1.0
 * @date ：Created in 2022/3/12 17:00
 * @description：
 */
public class MySqlDriverAgent implements ClassFileTransformer {

    private static final String TARGET_CLASS = "com.mysql.cj.jdbc.NonRegisteringDriver";
    private static final String TARGET_METHOD = "connect";
    private static final String TARGET_METHOD_DESC = "(Ljava/lang/String;Ljava/util/Properties;)Ljava/sql/Connection;";


    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        if (!TARGET_CLASS.replaceAll("\\.", "/").equals(className)) {
            return null;
        }
        try {
            return buildClass(loader);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    private byte[] buildClass(ClassLoader loader) throws Exception {
        ClassPool pool = new ClassPool();
        pool.insertClassPath(new LoaderClassPath(loader));
        CtClass ctClass = pool.get(TARGET_CLASS);
        CtMethod oldMethod = ctClass.getMethod(TARGET_METHOD, TARGET_METHOD_DESC);
        CtMethod newMethod = CtNewMethod.copy(oldMethod, ctClass, null);
        oldMethod.setName(oldMethod.getName() + "$agent");
        String endSrc = "result= javaagent.jdbc.JdbcCollects.proxyConnection((java.sql.Connection)result);";
        newMethod.setBody(String.format(source, "", TARGET_METHOD, "", endSrc));
        ctClass.addMethod(newMethod);
        return ctClass.toBytecode();
    }

    final static String source = "{\n"
            + "%s"
            + "        Object result=null;\n"
            + "       try {\n"
            + "            result=($w)%s$agent($$);\n"
            + "        } catch (Throwable e) {\n"
            + "%s"
            + "            throw e;\n"
            + "        }finally{\n"
            + "%s"
            + "        }\n"
            + "        return ($r) result;\n"
            + "}\n";
}
